# based on Stable range edge functions

# Initialises a matrix with n individuals
# individuals have a location (spX, spY) and are 
# initilised with gene for p.disp
# based on Stable range edge functions
init.inds<-function(n, spX, spY, hvar=FALSE, dval=99){
	X<-runif(n, 1, spX+1)%/%1 #grid ref
	Y<-runif(n, 1, spY+1)%/%1
	if (dval<=1) P<-rep(dval, n)
	else P<-runif(n, 0, 1) # dispersal probability
	if (hvar==FALSE) H<-rep(1, n)
	else H<-runif(n, 0, 1) #Habitat match
	cbind(X, Y, P, H)		
}

# initialises a habitat matrix based on homogenous area (maxX-nsteps, spY) and nsteps
# nsteps defines the rate at which habitat declines beyond homogenous area and >0
# homogenous area has default survival value of 1
# gradient ends five steps from the RHS of the space
init.hab<-function(maxX, spY, nsteps, hab.val=1){
	if (nsteps>maxX) return(print("Error, maxX>nsteps"))
	if (nsteps==0) return(matrix(data=1, nrow=maxX, ncol=spY))
	ncol1<-maxX-nsteps-5
	simple<-matrix(hab.val, nrow=spY, ncol=ncol1)
	simple2<-matrix(0, nrow=spY, ncol=5)
	d<-seq(hab.val, 0, length.out=nsteps+1)[-1]
	d<-rep(d, times=spY)
	grade<-t(matrix(d, ncol=spY, nrow=nsteps))
	t(cbind(simple, grade, simple2)) #same format as density matrix
}

# hassel commins population growth where b=1 (contest competition)
# generates expected fecundity
hass.comm<-function(lambda, K, N){
	exp<-lambda/(1+(lambda-1)*N/K)
}

# reproduces individuals based on density and habitat sepcific fitness, kills parents
# kills dispersing offspring based on dispersal cost
# disperses individuals based on inherited P
repro.disp<-function(popmatrix, hab.matrix, spX, spY, K, lambda, d.cost, n.list, mutn, adap=FALSE, dmut=TRUE){
	density<-table(factor(popmatrix[,"X"], levels=1:spX), factor(popmatrix[,"Y"], levels=1:spY)) # adult density through space
	temp<-(popmatrix[,"Y"]-1)*spX+popmatrix[,"X"] #matrix indexes for density/habitat matrices
	dens<-density[temp] # densities returned to each individual
	hab<-hab.matrix[temp] # habitat returned to each individual
	if (adap==TRUE) hab<-1-abs(hab-popmatrix[,"H"]) #decline in survival with habitat fit
	#hab<-ifelse(hab<0, 0, hab)
	offs<-rpois(length(popmatrix[,1]), hab*hass.comm(lambda, K, dens)) # weaned offspring based on fitness and density
	#offs<-rbinom(length(offs), offs, hab) # survival of offspring based on habitat
	if (length(offs)==0) return(NULL) #catch extinctions
	inds<-1:length(offs) #vector indexes for offspring
	inds<-rep(inds, times=offs)
	popmatrix<-popmatrix[inds,1:4] #replace parents with offspring. Offspring inherit location and P from parents
	if (class(popmatrix)=="numeric") popmatrix<-matrix(popmatrix, ncol=4, dimnames=list(NULL, c("X", "Y", "P", "H"))) #cases where only one individual left
	mut.sample<-rbinom(length(inds), 1, mutn) # individuals copping a mutation
	inds<-which(mut.sample==1)
	if (length(inds)>0 & dmut==TRUE){
		mutation<-rnorm(length(inds), popmatrix[inds, "P"], sd=0.05)
		mutation<-ifelse(mutation>1 | mutation<0, popmatrix[inds, "P"], mutation) #discard mutations outside bounds
		popmatrix[inds, "P"]<-mutation
	}
	mut.sample<-rbinom(length(inds), 1, mutn) # individuals copping a mutation
	inds<-which(mut.sample==1)
	if (length(inds)>0){
		mutation<-rnorm(length(inds), popmatrix[inds, "H"], sd=0.05)
		mutation<-ifelse(mutation<0 | mutation>1, popmatrix[inds, "H"], mutation) #discard mutations outside bounds
		popmatrix[inds, "H"]<-mutation
	}
	disp<-rbinom(length(popmatrix[,1]), 1, popmatrix[,"P"]) #disperse or not
	popmatrix<-cbind(popmatrix, disp)
	mort<-d.cost*disp #probability of dying during dispersal
	popmatrix<-subset(popmatrix, rbinom(length(mort), 1, mort)!=1) # remove those dying during dispersal
	popmatrix<-disperse(popmatrix, n.list, spX)
	popmatrix
}

#calculates neighbours on a grid excluding nonsense neighbours and returns a list of matrices of grid refs
#doesn't include source cell
# space is cylindrical
neighbours.init<-function(spX, spY){
	neigh<-vector("list", length=spX*spY)
	for (x in 1:spX){
		for (y in 1:spY){
			out<-matrix(nrow=8, ncol=2)
			ifelse((x-1)==0, X<-NA, X<-x-1) #correct for indices going to zero
			ifelse((y-1)==0, Y<-spY, Y<-y-1)
			ifelse((y+1)>spY, Yp<-1, Yp<-y+1) # correct for indices going to > space
			ifelse((x+1)>spX, Xp<-NA, Xp<-x+1)
			# Assign neighbours
			out[1,]<-c(Xp, y)
			out[2,]<-c(X, y)
			out[3,]<-c(x, Y)
			out[4,]<-c(Xp, Y)	
			out[5,]<-c(X, Y)
			out[6,]<-c(x, Yp)
			out[7,]<-c(Xp, Yp)
			out[8,]<-c(X, Yp)
			out<-subset(out, !is.na(apply(out, 1, sum)))
			neigh[[(y-1)*spX+x]]<-out
		}	
	}
	neigh	#return the list
}

# samples a matrix and returns one row of the matrix
matrix.sample<-function(mtrx, size=1){
	mtrx[sample(seq(length(mtrx[,1])), size=size),]	
}

# executes nearest neighbour dispersal where individuals disperse to neighbouring cell with probability=prob.d 
# is called after reproduction/survival
disperse<-function(popmatrix, n.list, spX){
	if (length(popmatrix[,3])==0) return(popmatrix)
	disp<-which(popmatrix[,"disp"]==1) #get indices of surviving dispersers
	if (length(disp)==0) return(popmatrix)
	sub.list<-n.list[(popmatrix[disp,"Y"]-1)*spX+popmatrix[disp,"X"]] #generate list of neighbour matrices
	sub.list<-lapply(sub.list, matrix.sample) #samples one row from each neighbour matrix
	sub.list<-do.call("rbind", sub.list) #places result into a matrix
	popmatrix[disp, 1:2]<-sub.list #assigns new locations back to popmatrix
	popmatrix
}

# plots population density and mean trait values through space
plotter<-function(popmatrix, spX, spY, K){
	density<-table(factor(popmatrix[,"X"], levels=1:spX), factor(popmatrix[,"Y"], levels=1:spY)) # adult density through space
	density<-density/K
	dispersal<-tapply(popmatrix[,"P"], list(factor(popmatrix[,"X"], levels=1:spX), factor(popmatrix[,"Y"], levels=1:spY)), mean) #mean trait values through space
	hval<-tapply(popmatrix[,"H"], list(factor(popmatrix[,"X"], levels=1:spX), factor(popmatrix[,"Y"], levels=1:spY)), mean) #mean trait values through space
	par(mfrow=c(3,1))
	image(density, zlim=c(0,2), col=(grey((12:0)/12)), main="Population density")
	image(dispersal, zlim=c(0,1), col=(grey((12:0)/12)), main="Dispersal values")
	image(hval, zlim=c(0,1), col=(grey((12:0)/12)), main="Habitat values")
}

col.sample<-function(pop, column){
	rws<-which(pop[,"X"]==column) 
	mean.d<-mean(pop[rws, "P"])
	var.d<-var(pop[rws,"P"])
	mean.h<-mean(pop[rws, "H"])
	var.h<-var(pop[rws,"H"])	
	out<-c(mean.d, var.d, mean.h, var.h)
	out
}

#mother function
mother<-function(n=20*40*10, ngens.init=100, gens.breach=1000, spY=20, steps, K=40, lambda, d.cost, mutn=0.05, plot=F, dval=99, spread=0){
	pop<-init.inds(n, 10, spY, hvar=T, dval)
	if (dval<=1) dmut<-FALSE
	else dmut<-TRUE
	maxX<-10+spread+steps+5
	hab<-init.hab(10, spY, 0)	
	nbrs<-neighbours.init(10, spY) #reveal only x up to 10
	for (i in 1:ngens.init){
		pop<-repro.disp(pop, hab, 10, spY, K, lambda, d.cost, nbrs, mutn, adap=T, dmut=dmut)
		if (plot==T) plotter(pop, maxX, spY, K)
	}
	init.traits<-col.sample(pop, 5) #sample initial means and variances
	nbrs<-neighbours.init(maxX, spY) #reveal all neighbours
	ngens<-ngens.init+spread+gens.breach
	hab<-init.hab(maxX, spY, steps)
	st.time<-0
	for (i in (ngens.init+1):ngens){
		pop<-repro.disp(pop, hab, maxX, spY, K, lambda, d.cost, nbrs, mutn, adap=T, dmut=dmut)
		if (plot==T) plotter(pop, maxX, spY, K)
		if (st.time==0 & sum((pop[,"X"])==(10+spread))>0) st.time<-i
		if ((st.time+1)==i){ #get 400 individuals closest to edge
			temp<-pop[rev(order(pop[,"X"]))[1:800], ]
			temp[,"X"]<-rep(1, 800)
			post.spread.traits<-col.sample(temp, 1)
		}
		if (max(pop[,"X"])==maxX | i==ngens) {
			post.breach.traits<-col.sample(pop, (10+spread))
			b.time<-ifelse(i==ngens, NA, i)
			out<-c(st.time, b.time, init.traits, post.spread.traits, post.breach.traits)
			names(out)<-c("start", "breach", "mean.d1", "var.d1", "mean.h1", "var.h1", "mean.d2", "var.d2", "mean.h2", "var.h2", "mean.d3", "var.d3", "mean.h3", "var.h3")
			return(out)
		}
	}
}

